import numpy as np
import logging
from dataclasses import dataclass
from typing import List, Tuple, Dict, Optional
import pandas as pd
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

@dataclass
class ProblemParams:
    """问题参数类，存储所有固定物理参数"""
    # 假目标位置 (x, y, z)
    fake_target: Tuple[float, float, float] = (0, 0, 0)
    # 真目标中心位置 (x, y, z)
    real_target: Tuple[float, float, float] = (0, 200, 5)
    # 导弹初始位置 (x, y, z)
    missile_init_pos: Tuple[float, float, float] = (20000, 0, 2000)
    # 导弹速度 (m/s)
    missile_speed: float = 300
    # 无人机初始位置 (x, y, z)
    uav_init_pos: Tuple[float, float, float] = (17800, 0, 1800)
    # 无人机速度范围 (m/s)
    uav_speed_range: Tuple[float, float] = (70, 140)
    # 烟幕有效半径 (m)
    smoke_radius: float = 10
    # 烟幕下沉速度 (m/s)
    smoke_sink_speed: float = 3
    # 烟幕有效时间 (s)
    smoke_effective_time: float = 20
    # 重力加速度 (m/s^2)
    gravity: float = 9.8
    # 时间步长 (s)
    time_step: float = 0.1
    
    def __post_init__(self):
        """计算导弹飞行轨迹的参数"""
        # 导弹到假目标的总距离
        self.missile_total_distance = np.sqrt(
            (self.missile_init_pos[0] - self.fake_target[0])**2 +
            (self.missile_init_pos[1] - self.fake_target[1])**2 +
            (self.missile_init_pos[2] - self.fake_target[2])**2
        )
        # 导弹总飞行时间
        self.missile_total_time = self.missile_total_distance / self.missile_speed
        # 导弹飞行方向单位向量
        self.missile_dir = np.array([
            (self.fake_target[0] - self.missile_init_pos[0]) / self.missile_total_distance,
            (self.fake_target[1] - self.missile_init_pos[1]) / self.missile_total_distance,
            (self.fake_target[2] - self.missile_init_pos[2]) / self.missile_total_distance
        ])
        logger.info(f"导弹总飞行时间: {self.missile_total_time:.2f}s, 总距离: {self.missile_total_distance:.2f}m")


class SmokeScreenResultGenerator:
    """烟幕弹投放结果生成器，用于生成表格所需数据"""
    
    def __init__(self, problem_params: ProblemParams = None):
        """初始化结果生成器"""
        self.problem_params = problem_params or ProblemParams()
        # 预计算导弹轨迹
        self.missile_trajectory = self._precompute_missile_trajectory()
    
    def _precompute_missile_trajectory(self) -> Dict[float, np.ndarray]:
        """预计算导弹轨迹以提高性能"""
        trajectory = {}
        time_points = np.arange(0, self.problem_params.missile_total_time + self.problem_params.time_step, 
                               self.problem_params.time_step)
        
        for t in time_points:
            trajectory[t] = self._missile_position(t)
            
        return trajectory
    
    def _missile_position(self, t: float) -> np.ndarray:
        """计算导弹在时刻t的位置"""
        if t < 0:
            return np.array(self.problem_params.missile_init_pos)
        if t > self.problem_params.missile_total_time:
            return np.array(self.problem_params.fake_target)
        
        distance_traveled = self.problem_params.missile_speed * t
        return np.array(self.problem_params.missile_init_pos) + distance_traveled * self.problem_params.missile_dir
    
    def calculate_positions(self, best_position: np.ndarray) -> Dict:
        """
        根据优化得到的最佳参数计算表格所需的所有位置和时间数据
        
        Args:
            best_position: 优化得到的最佳参数向量，顺序为：
                          [theta1, vf1, t_drop1, dt_int1, dt_int2, dt1, dt2, dt3]
            
        Returns:
            包含所有表格数据的字典
        """
        # 解析最佳参数
        theta1, vf1, t_drop1, dt_int1, dt_int2, dt1, dt2, dt3 = best_position
        
        # 计算第二、三枚烟幕的投放时刻
        t_drop2 = t_drop1 + dt_int1
        t_drop3 = t_drop2 + dt_int2
        
        # 计算各枚烟幕的投放位置
        drop_pos1 = self.smoke_drop_position(theta1, vf1, t_drop1)
        drop_pos2 = self.smoke_drop_position(theta1, vf1, t_drop2)
        drop_pos3 = self.smoke_drop_position(theta1, vf1, t_drop3)
        
        # 计算各枚烟幕的起爆位置
        det_pos1 = self.smoke_detonation_position(drop_pos1, dt1)
        det_pos2 = self.smoke_detonation_position(drop_pos2, dt2)
        det_pos3 = self.smoke_detonation_position(drop_pos3, dt3)
        
        # 计算总有效干扰时间
        total_effective_time = self.calculate_total_effective_time(best_position)
        
        # 转换无人机方向角为度（0-360度，逆时针为正）
        uav_direction = np.degrees(theta1) % 360
        
        # 组织结果数据
        result = {
            "无人机运动方向(度)": round(uav_direction, 2),
            "无人机运动速度(m/s)": round(vf1, 2),
            # 烟幕1数据
            "烟幕1投放时刻(s)": round(t_drop1, 2),
            "烟幕1投放点x坐标(m)": round(drop_pos1[0], 2),
            "烟幕1投放点y坐标(m)": round(drop_pos1[1], 2),
            "烟幕1投放点z坐标(m)": round(drop_pos1[2], 2),
            "烟幕1起爆时刻(s)": round(t_drop1 + dt1, 2),
            "烟幕1起爆点x坐标(m)": round(det_pos1[0], 2),
            "烟幕1起爆点y坐标(m)": round(det_pos1[1], 2),
            "烟幕1起爆点z坐标(m)": round(det_pos1[2], 2),
            # 烟幕2数据
            "烟幕2投放时刻(s)": round(t_drop2, 2),
            "烟幕2投放点x坐标(m)": round(drop_pos2[0], 2),
            "烟幕2投放点y坐标(m)": round(drop_pos2[1], 2),
            "烟幕2投放点z坐标(m)": round(drop_pos2[2], 2),
            "烟幕2起爆时刻(s)": round(t_drop2 + dt2, 2),
            "烟幕2起爆点x坐标(m)": round(det_pos2[0], 2),
            "烟幕2起爆点y坐标(m)": round(det_pos2[1], 2),
            "烟幕2起爆点z坐标(m)": round(det_pos2[2], 2),
            # 烟幕3数据
            "烟幕3投放时刻(s)": round(t_drop3, 2),
            "烟幕3投放点x坐标(m)": round(drop_pos3[0], 2),
            "烟幕3投放点y坐标(m)": round(drop_pos3[1], 2),
            "烟幕3投放点z坐标(m)": round(drop_pos3[2], 2),
            "烟幕3起爆时刻(s)": round(t_drop3 + dt3, 2),
            "烟幕3起爆点x坐标(m)": round(det_pos3[0], 2),
            "烟幕3起爆点y坐标(m)": round(det_pos3[1], 2),
            "烟幕3起爆点z坐标(m)": round(det_pos3[2], 2),
            # 总有效干扰时间
            "总有效干扰时间(s)": round(total_effective_time, 2)
        }
        
        return result
    
    def smoke_drop_position(self, theta: float, vf: float, t_drop: float) -> np.ndarray:
        """计算烟幕弹投放位置"""
        x = self.problem_params.uav_init_pos[0] + vf * np.cos(theta) * t_drop
        y = self.problem_params.uav_init_pos[1] + vf * np.sin(theta) * t_drop
        z = self.problem_params.uav_init_pos[2]  # 无人机保持水平飞行
        return np.array([x, y, z])
    
    def smoke_detonation_position(self, drop_pos: np.ndarray, dt: float) -> np.ndarray:
        """计算烟幕弹起爆位置"""
        # 水平位置不变，等于投放位置的水平坐标
        x = drop_pos[0]
        y = drop_pos[1]
        
        # 竖直位置：投放点z坐标减去自由落体距离
        z = drop_pos[2] - 0.5 * self.problem_params.gravity * dt**2
        
        return np.array([x, y, max(z, 0)])  # 确保z坐标不小于0（地面）
    
    def segment_sphere_intersection(self, A: np.ndarray, B: np.ndarray, C: np.ndarray, r: float) -> bool:
        """判断线段AB与球体(球心C，半径r)是否相交"""
        # 向量AB
        AB = B - A
        # 向量AC
        AC = C - A
        
        # 计算二次函数系数
        a = np.dot(AB, AB)
        b = 2 * np.dot(AC, AB)
        c = np.dot(AC, AC) - r**2
        
        # 计算判别式
        discriminant = b**2 - 4 * a * c
        
        # 如果判别式小于0，无交点
        if discriminant < 0:
            return False
        
        # 计算根
        sqrt_d = np.sqrt(discriminant)
        t1 = (-b - sqrt_d) / (2 * a)
        t2 = (-b + sqrt_d) / (2 * a)
        
        # 检查是否有根在[0, 1]区间内
        return (0 <= t1 <= 1) or (0 <= t2 <= 1)
    
    def single_smoke_effective_time(self, theta: float, vf: float, t_drop: float, dt: float) -> float:
        """计算单枚烟幕的有效遮蔽时间"""
        # 计算投放位置
        drop_pos = self.smoke_drop_position(theta, vf, t_drop)
        # 计算起爆时刻和失效时刻
        t_det = t_drop + dt
        t_start = max(t_det, 0)
        t_end = min(t_det + self.problem_params.smoke_effective_time, self.problem_params.missile_total_time)
        
        # 如果有效时间窗口不存在，返回0
        if t_start >= t_end:
            return 0.0
        
        # 使用预计算的导弹轨迹
        effective_time = 0.0
        time_points = np.arange(t_start, t_end + self.problem_params.time_step, self.problem_params.time_step)
        
        for t in time_points:
            if t > self.problem_params.missile_total_time:
                break
                
            # 获取导弹位置
            missile_pos = self.missile_trajectory[round(t / self.problem_params.time_step) * self.problem_params.time_step]
            
            # 计算烟幕球心位置
            smoke_center = self.smoke_detonation_position(drop_pos, dt)
            # 考虑烟幕下沉
            sink_distance = self.problem_params.smoke_sink_speed * (t - t_det)
            smoke_center[2] = max(smoke_center[2] - sink_distance, 0)  # 确保不低于地面
            
            # 判断线段(导弹-真目标)与烟幕球体是否相交
            if self.segment_sphere_intersection(
                missile_pos, 
                np.array(self.problem_params.real_target), 
                smoke_center, 
                self.problem_params.smoke_radius
            ):
                effective_time += self.problem_params.time_step
        
        return effective_time
    
    def calculate_total_effective_time(self, best_position: np.ndarray) -> float:
        """计算总有效干扰时间"""
        theta1, vf1, t_drop1, dt_int1, dt_int2, dt1, dt2, dt3 = best_position
        
        # 计算第二、三枚烟幕的投放时刻
        t_drop2 = t_drop1 + dt_int1
        t_drop3 = t_drop2 + dt_int2
        
        # 计算每枚烟幕的有效遮蔽时间
        t1 = self.single_smoke_effective_time(theta1, vf1, t_drop1, dt1)
        t2 = self.single_smoke_effective_time(theta1, vf1, t_drop2, dt2)
        t3 = self.single_smoke_effective_time(theta1, vf1, t_drop3, dt3)
        
        # 返回总有效遮蔽时间
        return t1 + t2 + t3
    
    def generate_table(self, best_position: np.ndarray) -> pd.DataFrame:
        """生成表格形式的结果"""
        result = self.calculate_positions(best_position)
        
        # 组织表格数据
        table_data = {
            "参数": [
                "无人机运动方向(度)",
                "无人机运动速度(m/s)",
                "总有效干扰时间(s)",
                "",
                "烟幕1投放时刻(s)",
                "烟幕1投放点x坐标(m)",
                "烟幕1投放点y坐标(m)",
                "烟幕1投放点z坐标(m)",
                "烟幕1起爆时刻(s)",
                "烟幕1起爆点x坐标(m)",
                "烟幕1起爆点y坐标(m)",
                "烟幕1起爆点z坐标(m)",
                "",
                "烟幕2投放时刻(s)",
                "烟幕2投放点x坐标(m)",
                "烟幕2投放点y坐标(m)",
                "烟幕2投放点z坐标(m)",
                "烟幕2起爆时刻(s)",
                "烟幕2起爆点x坐标(m)",
                "烟幕2起爆点y坐标(m)",
                "烟幕2起爆点z坐标(m)",
                "",
                "烟幕3投放时刻(s)",
                "烟幕3投放点x坐标(m)",
                "烟幕3投放点y坐标(m)",
                "烟幕3投放点z坐标(m)",
                "烟幕3起爆时刻(s)",
                "烟幕3起爆点x坐标(m)",
                "烟幕3起爆点y坐标(m)",
                "烟幕3起爆点z坐标(m)"
            ],
            "值": [
                result["无人机运动方向(度)"],
                result["无人机运动速度(m/s)"],
                result["总有效干扰时间(s)"],
                "",
                result["烟幕1投放时刻(s)"],
                result["烟幕1投放点x坐标(m)"],
                result["烟幕1投放点y坐标(m)"],
                result["烟幕1投放点z坐标(m)"],
                result["烟幕1起爆时刻(s)"],
                result["烟幕1起爆点x坐标(m)"],
                result["烟幕1起爆点y坐标(m)"],
                result["烟幕1起爆点z坐标(m)"],
                "",
                result["烟幕2投放时刻(s)"],
                result["烟幕2投放点x坐标(m)"],
                result["烟幕2投放点y坐标(m)"],
                result["烟幕2投放点z坐标(m)"],
                result["烟幕2起爆时刻(s)"],
                result["烟幕2起爆点x坐标(m)"],
                result["烟幕2起爆点y坐标(m)"],
                result["烟幕2起爆点z坐标(m)"],
                "",
                result["烟幕3投放时刻(s)"],
                result["烟幕3投放点x坐标(m)"],
                result["烟幕3投放点y坐标(m)"],
                result["烟幕3投放点z坐标(m)"],
                result["烟幕3起爆时刻(s)"],
                result["烟幕3起爆点x坐标(m)"],
                result["烟幕3起爆点y坐标(m)"],
                result["烟幕3起爆点z坐标(m)"]
            ]
        }
        
        # 创建DataFrame
        df = pd.DataFrame(table_data)
        
        return df
    
    def visualize_trajectories(self, best_position: np.ndarray):
        """可视化导弹轨迹和烟幕投放点"""
        fig = plt.figure(figsize=(12, 10))
        ax = fig.add_subplot(111, projection='3d')
        
        # 绘制导弹轨迹
        times = list(self.missile_trajectory.keys())
        missile_positions = list(self.missile_trajectory.values())
        missile_x = [pos[0] for pos in missile_positions]
        missile_y = [pos[1] for pos in missile_positions]
        missile_z = [pos[2] for pos in missile_positions]
        
        ax.plot(missile_x, missile_y, missile_z, 'b-', label='导弹轨迹', linewidth=2)
        
        # 绘制假目标和真目标
        ax.scatter(*self.problem_params.fake_target, c='r', marker='o', s=100, label='假目标')
        ax.scatter(*self.problem_params.real_target, c='g', marker='^', s=100, label='真目标')
        
        # 绘制无人机初始位置
        ax.scatter(*self.problem_params.uav_init_pos, c='k', marker='s', s=100, label='无人机初始位置')
        
        # 绘制烟幕投放点和起爆点
        result = self.calculate_positions(best_position)
        
        for i in range(1, 4):
            drop_x = result[f"烟幕{i}投放点x坐标(m)"]
            drop_y = result[f"烟幕{i}投放点y坐标(m)"]
            drop_z = result[f"烟幕{i}投放点z坐标(m)"]
            
            det_x = result[f"烟幕{i}起爆点x坐标(m)"]
            det_y = result[f"烟幕{i}起爆点y坐标(m)"]
            det_z = result[f"烟幕{i}起爆点z坐标(m)"]
            
            ax.scatter(drop_x, drop_y, drop_z, c='m', marker='o', s=80, label=f'烟幕{i}投放点')
            ax.scatter(det_x, det_y, det_z, c='c', marker='x', s=80, label=f'烟幕{i}起爆点')
            
            # 绘制投放点到起爆点的连线
            ax.plot([drop_x, det_x], [drop_y, det_y], [drop_z, det_z], 'm--', linewidth=1)
        
        ax.set_xlabel('X (m)')
        ax.set_ylabel('Y (m)')
        ax.set_zlabel('Z (m)')
        ax.set_title('导弹轨迹与烟幕投放可视化')
        
        # 限制图例项数量，避免重复
        handles, labels = ax.get_legend_handles_labels()
        by_label = dict(zip(labels, handles))
        ax.legend(by_label.values(), by_label.keys(), loc='center left', bbox_to_anchor=(1, 0.5))
        
        plt.tight_layout()
        plt.show()
    
    def save_results(self, best_position: np.ndarray, filename: str = "smoke_screen_results.xlsx"):
        """将结果保存到Excel文件"""
        try:
            df = self.generate_table(best_position)
            df.to_excel(filename, index=False)
            logger.info(f"结果已成功保存到 {filename}")
            
            # 同时生成可视化图表
            self.visualize_trajectories(best_position)
            plt.savefig(filename.replace('.xlsx', '_visualization.png'), dpi=300, bbox_inches='tight')
            logger.info(f"可视化图表已保存到 {filename.replace('.xlsx', '_visualization.png')}")
            
            return True
        except Exception as e:
            logger.error(f"保存结果时出错: {str(e)}")
            return False


# 使用示例
def main():
    # 假设这是从优化算法得到的最佳参数
    # 实际使用时应替换为真实的优化结果
    best_position = np.array([
        3.5,    # theta1 (无人机方向角，弧度)
        100,    # vf1 (无人机速度，m/s)
        15,     # t_drop1 (第一枚投放时间，s)
        3,      # dt_int1 (第一与第二枚间隔，s)
        4,      # dt_int2 (第二与第三枚间隔，s)
        12,     # dt1 (第一枚起爆延迟，s)
        11,     # dt2 (第二枚起爆延迟，s)
        10      # dt3 (第三枚起爆延迟，s)
    ])
    
    # 创建结果生成器
    generator = SmokeScreenResultGenerator()
    
    # 生成并打印表格
    result_table = generator.generate_table(best_position)
    print("烟幕弹投放策略优化结果:")
    print(result_table.to_string(index=False))
    
    # 保存结果到Excel
    generator.save_results(best_position)

if __name__ == "__main__":
    main()